# Copyright (C) 2025 ByteDance and/or its affiliates

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

FORCEFIELD = {"FF14SB": "leaprc.protein.ff14SB"}
FORCEFIELD_WATER_ION = {"FF14SB": "leaprc.water.tip3p"}

ION_RESNAME_TO_ATOMNAME = {
    "AG": "AG",
    "AL": "AL",
    "Ag": "Ag",
    "BA": "BA",
    "BR": "BR",
    "Be": "Be",
    "CA": "CA",
    "CD": "CD",
    "CE": "CE",
    "CL": "CL",
    "CO": "CO",
    "CR": "CR",
    "CS": "CS",
    "CU": "CU",
    "CU1": "CU",
    "Ce": "Ce",
    "Cl-": "Cl-",
    "Cr": "Cr",
    "Dy": "Dy",
    "EU": "EU",
    "EU3": "EU3",
    "Er": "Er",
    "F": "F",
    "FE": "FE",
    "FE2": "FE2",
    "GD3": "GD",
    "H3O+": "O",
    "HE+": "H",
    "HG": "HG",
    "HZ+": "H",
    "Hf": "Hf",
    "IN": "IN",
    "IOD": "I",
    "K": "K",
    "K+": "K+",
    "LA": "LA",
    "LI": "LI",
    "LU": "LU",
    "MG": "MG",
    "MN": "MN",
    "NA": "NA",
    "NH4": "N",
    "NI": "NI",
    "Na+": "Na+",
    "Nd": "Nd",
    "PB": "PB",
    "PD": "PD",
    "PR": "PR",
    "PT": "PT",
    "Pu": "Pu",
    "RB": "RB",
    "Ra": "Ra",
    "SM": "SM",
    "SR": "SR",
    "Sm": "Sm",
    "Sn": "Sn",
    "TB": "TB",
    "TL": "TL",
    "Th": "Th",
    "Tl": "Tl",
    "Tm": "Tm",
    "U4+": "U",
    "V2+": "V2+",
    "Y": "Y",
    "YB2": "YB2",
    "ZN": "ZN",
    "Zr": "Zr",
}

ION_RESNAME_TO_ATOMTYPE = {
    "AG": "Ag+",
    "AL": "Al3+",
    "Ag": "Ag2+",
    "BA": "Ba2+",
    "BR": "Br-",
    "Be": "Be2+",
    "CA": "Ca2+",
    "CD": "Cd2+",
    "CE": "Ce3+",
    "CL": "Cl-",
    "CO": "Co2+",
    "CR": "Cr3+",
    "CS": "Cs+",
    "CU": "Cu2+",
    "CU1": "Cu+",
    "Ce": "Ce4+",
    "Cl-": "Cl-",
    "Cr": "Cr2+",
    "Dy": "Dy3+",
    "EU": "Eu2+",
    "EU3": "Eu3+",
    "Er": "Er3+",
    "F": "F-",
    "FE": "Fe3+",
    "FE2": "Fe2+",
    "GD3": "Gd3+",
    "H3O+": "H3O+",
    "HE+": "HE+",
    "HG": "Hg2+",
    "HZ+": "HZ+",
    "Hf": "Hf4+",
    "IN": "In3+",
    "IOD": "I-",
    "K": "K+",
    "K+": "K+",
    "LA": "La3+",
    "LI": "Li+",
    "LU": "Lu3+",
    "MG": "Mg2+",
    "MN": "Mn2+",
    "NA": "Na+",
    "NH4": "NH4+",
    "NI": "Ni2+",
    "Na+": "Na+",
    "Nd": "Nd3+",
    "PB": "Pb2+",
    "PD": "Pd2+",
    "PR": "Pr3+",
    "PT": "Pt2+",
    "Pu": "Pu4+",
    "RB": "Rb+",
    "Ra": "Ra2+",
    "SM": "Sm3+",
    "SR": "Sr2+",
    "Sm": "Sm2+",
    "Sn": "Sn2+",
    "TB": "Tb3+",
    "TL": "Tl+",
    "Th": "Th4+",
    "Tl": "Tl3+",
    "Tm": "Tm3+",
    "U4+": "U4+",
    "V2+": "V2+",
    "Y": "Y3+",
    "YB2": "Yb2+",
    "ZN": "Zn2+",
    "Zr": "Zr4+",
}

AMBERFF_ATOMTYPES = [
    "H",
    "HO",
    "HS",
    "HC",
    "H1",
    "H2",
    "H3",
    "HP",
    "HA",
    "H4",
    "H5",
    "HZ",
    "O",
    "O2",
    "OH",
    "OS",
    "OP",
    "C*",
    "CA",
    "CB",
    "CC",
    "CD",
    "CK",
    "CM",
    "CN",
    "CQ",
    "CR",
    "CV",
    "CW",
    "CY",
    "CZ",
    "CP",
    "CS",
    "CI",
    "C5",
    "C4",
    "CT",
    "CX",
    "C",
    "N",
    "NA",
    "N2",
    "N*",
    "NC",
    "NB",
    "NT",
    "NY",
    "N3",
    "S",
    "SH",
    "P",
    "MG",
    "C0",
    "F",
    "Cl",
    "Br",
    "I",
    "2C",
    "3C",
    "C8",
    "CO",
    "FE",
    "LO",
    "LC",
    "NP",
    "NO",
]
GAFF2_ATOMTYPES = [
    "c",
    "cs",
    "c1",
    "c2",
    "c3",
    "ca",
    "cp",
    "cq",
    "cc",
    "cd",
    "ce",
    "cf",
    "cg",
    "ch",
    "cx",
    "cy",
    "c5",
    "c6",
    "cu",
    "cv",
    "cz",
    "h1",
    "h2",
    "h3",
    "h4",
    "h5",
    "ha",
    "hc",
    "hn",
    "ho",
    "hp",
    "hs",
    "hw",
    "hx",
    "f",
    "cl",
    "br",
    "i",
    "n",
    "n1",
    "n2",
    "n3",
    "n4",
    "na",
    "nb",
    "nc",
    "nd",
    "ne",
    "nf",
    "nh",
    "no",
    "ns",
    "nt",
    "nx",
    "ny",
    "nz",
    "n+",
    "nu",
    "nv",
    "n7",
    "n8",
    "n9",
    "ni",
    "nj",
    "nk",
    "nl",
    "nm",
    "nn",
    "np",
    "nq",
    "n5",
    "n6",
    "o",
    "oh",
    "op",
    "oq",
    "os",
    "ow",
    "p2",
    "p3",
    "p4",
    "p5",
    "pb",
    "pc",
    "pd",
    "pe",
    "pf",
    "px",
    "py",
    "s",
    "s2",
    "s4",
    "s6",
    "sh",
    "ss",
    "sx",
    "sy",
    "sp",
    "sq",
]


# WARNING: atom type K+ corresponds to two different resname and atomname
AMBERFF_ATOMTYPES = (
    GAFF2_ATOMTYPES + AMBERFF_ATOMTYPES + list(set(ION_RESNAME_TO_ATOMTYPE.values()))
)
